package com.gmrz.webauthn.protocol.v1.processor.attestation;

import co.nstant.in.cbor.model.DataItem;
import com.gmrz.service.challengestore.ChallengeStoreException;
import com.gmrz.uaf.common.Constants;
import com.gmrz.uaf.common.GuiceUtil;
import com.gmrz.uaf.crypto.CryptoEngine;
import com.gmrz.util.CryUtil;
import com.gmrz.webauthn.common.CBORDecodeUtil;
import com.gmrz.webauthn.common.WebAuthnConstants;
import com.gmrz.webauthn.protocol.v1.json.WebAuthnSchemaBuilder;
import com.gmrz.webauthn.protocol.v1.processor.WebAuthnCborException;
import com.gmrz.webauthn.protocol.v1.schema.*;
import com.gmrz.webauthn.request.WebAuthnCredentialContext;
import com.gmrz.util.Convert;
import com.gmrz.uaf.db.DAOException;
import com.gmrz.webauthn.protocol.v1.processor.WebAuthnErrorCode;
import com.gmrz.webauthn.protocol.v1.processor.WebAuthnException;

import java.nio.ByteBuffer;
import java.nio.charset.Charset;
import java.security.cert.Certificate;
import java.security.cert.X509Certificate;
import java.sql.SQLException;
import java.util.List;
import java.util.Map;

import org.apache.commons.codec.binary.Base64;
import org.apache.commons.codec.binary.Hex;
import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;

public class U2FAttestationProcessor
        extends WebAuthnBaseAssertionProcessor
        implements IAttestationProcessor
{
    private static final Logger LOG = LogManager.getLogger(U2FAttestationProcessor.class);
    public static final byte REGISTRATION_SIGNED_RESERVED_BYTE_VALUE = (byte) 0x00;

    public U2FAttestationProcessor(CryptoEngine cryptoEngine){
        this.cryptoEngine = cryptoEngine;
    }

    public WebAuthnConstants.AttestationResult verify(WebAuthnCredentialContext publicKeyCredential, WebAuthenticator authenticator,boolean strictMode)
            throws WebAuthnException, WebAuthnCborException, ChallengeStoreException, SQLException, DAOException {
        if ((publicKeyCredential == null) || (null == publicKeyCredential.getResponse() ))
        {
            String errorMsg = "public key credential is not present for attestation verification";
            LOG.error(errorMsg);
            return WebAuthnConstants.AttestationResult.UNAVAILABLE;
        }

        List<DataItem> dataitems = CBORDecodeUtil.cborDecode(Base64.decodeBase64(publicKeyCredential.getResponse().getAttestationObject()));
        co.nstant.in.cbor.model.Map cborMap = (co.nstant.in.cbor.model.Map) dataitems.get(0);

        java.util.Map<String, Object> map = CBORDecodeUtil.cborMapConversion(cborMap);

        byte[] rawAuthData = (byte[]) map.get("authData");

        AuthenticatorData authenticatorData = GuiceUtil.getProcessorInjector().getInstance(AuthenticatorData.class);
        authenticatorData.decode(rawAuthData);


        String formatName = (String) map.get("fmt");
        AttestationFormat format = AttestationFormat.forName(formatName);

        Object attStmtObj = map.get("attStmt");

        CredentialAttestation credentialAttestation = new CredentialAttestation().withAuthenticatorData(authenticatorData).withAuthenticatorDataBytes(rawAuthData).withAttestationStatement(attStmtObj).withFormat(format);
        credentialAttestation.setAuthenticatorDataBytes(rawAuthData);

        Map<String, Object> attStmt = (Map<String, Object>) credentialAttestation.getAttestationStatement();

        String clientDataStr = publicKeyCredential.getResponse().getClientDataJSON();
        String clientDataDecodedStr = new String(Base64.decodeBase64(clientDataStr), Charset.forName(Constants.UTF_8));
        CollectedClientData clientData = WebAuthnSchemaBuilder.getGson().fromJson(clientDataDecodedStr, CollectedClientData.class);

        byte[] signature = (byte[])attStmt.get("sig");

        List<Certificate> x509decodeCerts = extractCertChain(attStmt);

        X509Certificate attestationCertificate = (X509Certificate)x509decodeCerts.get(0);

        LOG.debug(" Attestation certificate: {}", attestationCertificate);
        try
        {
            byte[] signedBytes = generateSignedBytes(authenticatorData,Base64.decodeBase64(clientDataStr));
            LOG.debug("Verifying signature of bytes " + Hex.encodeHexString(signedBytes));
            boolean signatureVerification = cryptoEngine.verifySignature(2,attestationCertificate, signedBytes, signature);
            if (!signatureVerification) {
                throw new WebAuthnException(WebAuthnErrorCode.ERROR_ATTESTATION_SIGNATURE_VALIDATION_FAILED, "Signature is invalid");
            }
        }
        catch (Exception e)
        {
            String errorMessage = "Error while verifying signature";
            logAndUpdateResult(errorMessage, WebAuthnErrorCode.ERROR_ATTESTATION_SIGNATURE_VALIDATION_FAILED, e);
            return this.attResult;
        }
        LOG.debug("U2F Attestation signature verification is successful");

        return this.attResult;
    }

    private List<Certificate> extractCertChain(Map<String, Object> attStmt)
            throws WebAuthnException
    {
        List<Object> x509Certs = (List)attStmt.get("x5c");
        List<Certificate> x509decodeCerts = getDecodedX509Certs(x509Certs);
        return x509decodeCerts;
    }

    private byte[] generateSignedBytes(AuthenticatorData authData,byte[] clientDataBytes)
            throws WebAuthnException
    {
        byte[] applicationSha256 = null;
        byte[] challengeSha256 = null;
        byte[] keyHandle = null;
        byte[] userPublicKey = null;

        applicationSha256 = authData.getRpID();
        challengeSha256 = CryUtil.getSHA256(clientDataBytes);

        keyHandle = Convert.fromBase64(authData.getAttestation().getCredentialID());

        userPublicKey = Convert.fromBase64(authData.getAttestation().getCredentialPublicKey().getPublicKeyInRawFormat());

        byte[] signedBytes = encodeRegistrationSignedBytes(applicationSha256, challengeSha256, keyHandle, userPublicKey);

        return signedBytes;
    }

    public static byte[] encodeRegistrationSignedBytes(byte[] applicationSha256,
                                                       byte[] challengeSha256, byte[] keyHandle, byte[] userPublicKey) {
        byte[] signedData = new byte[1 + applicationSha256.length + challengeSha256.length
                + keyHandle.length + userPublicKey.length];
        ByteBuffer.wrap(signedData)
                .put(REGISTRATION_SIGNED_RESERVED_BYTE_VALUE) // RFU
                .put(applicationSha256)
                .put(challengeSha256)
                .put(keyHandle)
                .put(userPublicKey);
        return signedData;
    }

}

